# -*- coding: UTF-8 -*-
import re
import numpy as np
import pandas as pd

result = {}

def action(config, run_date):
    # 解析配置文件
    conf = config['akshare']['stocks']['CN']
    level = conf['level']
    symbols = conf['symbols']
    rules = conf['rules']

    file_path = 'data/akshare/stocks/CN'
    if level == '5M':
        process_5m(symbols, rules, file_path, run_date)


def process_5m(symbols, rules, file_path, run_date):
    for symbol in symbols:
        # 使用 pandas 来加载上一步获取到的交易数据
        data = pd.read_csv(f'{file_path}/{symbol}-{run_date}-5M.csv')
        for rule in rules:
            if rule['name'] == 'MMA':
                data, signal = process_mma(data)
            else:
                continue
            # 把满足条件的标的添加到结果集
            if not signal:
                continue
            if symbol in result:
                result[symbol].append(f'{rule["desc"]}: {signal}')
            else:
                result[symbol] = [f'{rule["desc"]}: {signal}']
        # 将包含指标列的数据保存为一个新文件
        data.to_csv(f'{file_path}/{symbol}-{run_date}-5M-indicators.csv', float_format='%.2f')
    if result:
        with open(f'{file_path}/result-{run_date}.csv', 'w') as f:
            for symbol, msg_list in result.items():
                f.write(f'{symbol}\t' + ', '.join(msg_list) + '\n')

def calculate_rsi(df, window=14):
    # 计算相对强弱指数（RSI）
    delta = df['收盘'].diff()
    gain = (delta.where(delta > 0, 0)).ewm(alpha=1/window, adjust=False).mean()
    loss = (-delta.where(delta < 0, 0)).ewm(alpha=1/window, adjust=False).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi

def check_signal(df, signal_col, cross_col, ma_col, max_bars=3):
    df[signal_col] = 0
    cross_indices = df.index[df[cross_col]]
    
    for idx in cross_indices:
        # 选取交叉及之后的 3 根 K 线窗口
        window = df.loc[idx: idx + max_bars]
        # 检查条件：价格与均线关系 + RSI区间
        if signal_col == 'b_signal':
            condition = (window['最低'] > window[ma_col]) & (50 <= window['RSI']) & (window['RSI'] <= 70)
        else:
            condition = (window[ma_col] > window['最高']) & (30 <= window['RSI']) & (window['RSI'] <= 50)
        # 找到首次满足条件的位置
        first_match = condition.idxmax() if condition.any() else None 
        if first_match is not None:
            df.loc[first_match, signal_col] = 1 
    return df
    

def process_mma(df):
    # 基础指标的计算
    # 先以收盘价计算 MA10 作为快线
    df['MA10'] = df['收盘'].rolling(window=10, min_periods=1).mean()
    # 再对 MA10 计算 EMA10 作为慢线
    df['EMA10'] = df['MA10'].ewm(span=10, adjust=False).mean()
    # 计算 RSI 作为后面的信号过滤条件
    df['RSI'] = calculate_rsi(df).rolling(window=5, min_periods=1).mean()

    # 策略规则的计算
    # 首先找到快慢线交叉的位置: 根据当前快慢线的位置关系以及上根K线两者的位置关系来判断是上穿还是下穿
    df['cross_up'] = (df['MA10'] > df['EMA10']) & (df['MA10'].shift(1) <= df['EMA10'].shift(1))
    df['cross_dw'] = (df['MA10'] < df['EMA10']) & (df['MA10'].shift(1) >= df['EMA10'].shift(1))

    # 然后在交叉信号之后的 3 根 K 线内判断是否触发了买卖信号
    df['b_signal'] = 0  
    df['s_signal'] = 0 
    df = check_signal(df, 'b_signal', 'cross_up', 'MA10')
    df = check_signal(df, 's_signal', 'cross_dw', 'MA10')

    return df, '买入' if df.iloc[-1]['b_signal'] else '卖出' if df.iloc[-1]['s_signal'] else ''


def main():
    symbols = ['000300']
    rules = [{'name': 'MMA', 'desc': '均线交叉+RSI过滤'}]
    file_path = '/tmp'
    run_date = '20241216'
    process_5m(symbols, rules, file_path, run_date)


if __name__ == '__main__':
    main()
    